package com.meizhao.api.filter;

import com.meizhao.api.utils.HttpUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.owasp.esapi.ESAPI;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.util.StreamUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import java.util.regex.Pattern;


public class SecurityHttpServletRequestWrapper extends HttpServletRequestWrapper {

    private final byte[] body;

    private Map<String, String[]> params;

    public SecurityHttpServletRequestWrapper(HttpServletRequest servletRequest) throws IOException{
        super(servletRequest);
        params = servletRequest.getParameterMap();
        if (servletRequest.getHeader(HttpHeaders.CONTENT_TYPE) != null && servletRequest.getHeader(HttpHeaders.CONTENT_TYPE).startsWith(MediaType.APPLICATION_FORM_URLENCODED_VALUE)) {
            body = null;
        } else {
            body = StreamUtils.copyToByteArray(servletRequest.getInputStream());
        }
    }

    @Override
    public String[] getParameterValues(String parameter) {

        String[] values = super.getParameterValues(parameter);
        if (values == null) {
            return null;
        }

        int count = values.length;
        String[] encodedValues = new String[count];
        for (int i = 0; i < count; i++) {
            encodedValues[i] = cleanXSS(values[i]);
        }
        return encodedValues;

    }


    @Override
    public String getHeader(String name) {

        String value = super.getHeader(name);
        if (value == null) {
            return null;
        }
        return cleanXSS(value);

    }


    private String cleanXSS(String value) {
        if (value != null) {
            // 推荐使用ESAPI库来避免脚本攻击
            value = ESAPI.encoder().canonicalize(value);

            // 避免空字符串
            value = value.trim();

            value = value.replaceAll("<", "& lt;").replaceAll(">", "& gt;");

            value = value.replaceAll("\\(", "& #40;").replaceAll("\\)", "& #41;");

            value = value.replaceAll("'", "& #39;");

            value = value.replaceAll("[\\\"\\\'][\\s]*javascript:(.*)[\\\"\\\']", "\"\"");

            // 避免script 标签
            Pattern scriptPattern = Pattern.compile("<script>(.*?)</script>", Pattern.CASE_INSENSITIVE);
            value = scriptPattern.matcher(value).replaceAll("");

            // 避免src形式的表达式
            scriptPattern = Pattern.compile("src[\r\n]*=[\r\n]*\\\'(.*?)\\\'", Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
            value = scriptPattern.matcher(value).replaceAll("");

            // 删除单个的 </script> 标签
            scriptPattern = Pattern.compile("</script>", Pattern.CASE_INSENSITIVE);
            value = scriptPattern.matcher(value).replaceAll("");

            // 删除单个的<script ...> 标签
            scriptPattern = Pattern.compile("<script(.*?)>", Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
            value = scriptPattern.matcher(value).replaceAll("");

            // 避免 eval(...) 形式表达式
            scriptPattern = Pattern.compile("eval\\((.*?)\\)", Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
            value = scriptPattern.matcher(value).replaceAll("");

            // 避免 e­xpression(...) 表达式
            scriptPattern = Pattern.compile("expression\\((.*?)\\)", Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
            value = scriptPattern.matcher(value).replaceAll("");

            // 避免 javascript: 表达式
            scriptPattern = Pattern.compile("javascript:", Pattern.CASE_INSENSITIVE);
            value = scriptPattern.matcher(value).replaceAll("");

            // 避免 vbscript: 表达式
            scriptPattern = Pattern.compile("vbscript:", Pattern.CASE_INSENSITIVE);
            value = scriptPattern.matcher(value).replaceAll("");

            // 避免 οnlοad= 表达式
            scriptPattern = Pattern.compile("onload(.*?)=", Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
            value = scriptPattern.matcher(value).replaceAll("");

            // 避免 onXX= 表达式
            scriptPattern = Pattern.compile("on.*(.*?)=", Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
            value = scriptPattern.matcher(value).replaceAll("");

        }
        return value;
    }

    private ServletInputStream wrapperInputStream() throws IOException {

        final ByteArrayInputStream bais = new ByteArrayInputStream(body);
        return new ServletInputStream() {

            @Override
            public int read() throws IOException {
                return bais.read();
            }

            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener listener) {

            }

        };
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {

        ServletInputStream stream = wrapperInputStream();

        //非json类型，直接返回
        if (!super.getHeader(HttpHeaders.CONTENT_TYPE).contains(MediaType.APPLICATION_JSON_VALUE)) {
            return stream;
        }
        String json = IOUtils.toString(stream, StandardCharsets.UTF_8);

        if (StringUtils.isBlank(json)) {
            return stream;
        }

        //System.out.println("转化前参数：" + json);
        Map<String, Object> map = modifyParams(json);
        //System.out.println("转化后参数：" + Utils.getObjectMapper().writeValueAsString(map));

        ByteArrayInputStream bis = new ByteArrayInputStream(HttpUtils.getObjectMapper().writeValueAsString(map).getBytes(StandardCharsets.UTF_8));

        return new NewServletInputStream(bis);
    }


    private static class NewServletInputStream extends ServletInputStream{

        private ByteArrayInputStream bis;

        public NewServletInputStream(ByteArrayInputStream bis) {
            this.bis = bis;
        }

        @Override
        public boolean isFinished() {
            return true;
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setReadListener(ReadListener readListener) {

        }

        @Override
        public int read() throws IOException {
            return bis.read();
        }
    }

    private Map<String, Object> modifyParams(String json) throws IOException{

        Map<String, Object> params2 = HttpUtils.getObjectMapper().readValue(json, Map.class);
        Map<String, Object> maps = new HashMap<>(params2.size());
        for (String key : params2.keySet()) {
            Object value = getValue(params2.get(key));
            maps.put(key, value);
        }
        return maps;
    }

    private Object getValue(Object obj) {

        if (obj == null) {
            return null;
        }
        String type = obj.getClass().getName();
        // 对字符串的处理
        if ("java.lang.String".equals(type)) {
            obj = (Object)(cleanXSS(obj.toString()));
        }
        return obj;
    }

    @Override
    public String getParameter(String name) {
        String[] values = params.get(name);
        if (values == null || values.length == 0) {
            return null;
        }
        return cleanXSS(values[0]);
    }
}